# -*-coding:utf-8-*-
import os
import glob
import cv2
from tracker import KCFTracker
import rospy
from sensor_msgs.msg import RegionOfInterest as ROI

cap = cv2.VideoCapture(0)

ret, img = cap.read()

print(img.shape)
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter("out.avi", fourcc, 24, (img.shape[1], img.shape[0]))
if not cap.isOpened():
    print("Could not open video");
    exit(-1)
tracker_types = ['BOOSTING', 'MIL', 'KCF', 'TLD', 'MEDIANFLOW', 'GOTURN', 'MOSSE']
tracker_type = tracker_types[2]
(major_ver, minor_ver, subminor_ver) = (cv2.__version__).split('.')
if int(minor_ver) < 3:
    tracker1 = cv2.Tracker_create(tracker_type)
else:
    if tracker_type == 'BOOSTING':
        tracker1 = cv2.TrackerBoosting_create()
    if tracker_type == 'MIL':
        tracker1 = cv2.TrackerMIL_create()
    if tracker_type == 'KCF':
        tracker1 = cv2.TrackerKCF_create()
    if tracker_type == 'TLD':
        tracker1 = cv2.TrackerTLD_create()
    if tracker_type == 'MEDIANFLOW':
        tracker1 = cv2.TrackerMedianFlow_create()
    if tracker_type == 'GOTURN':
        tracker1 = cv2.TrackerGOTURN_create()
    if tracker_type == 'MOSSE':
        tracker1 = cv2.TrackerGOTURN_create()

selection = None
track_window = None
drag_start = None


# 鼠标选框（做目标跟踪框）
def onmouse(event, x, y, flags, param):
    global selection, track_window, drag_start
    if event == cv2.EVENT_LBUTTONDOWN:
        drag_start = (x, y)
        track_window = None
    if drag_start:
        xmin = min(x, drag_start[0])
        ymin = min(y, drag_start[1])
        xmax = max(x, drag_start[0])
        ymax = max(y, drag_start[1])
        selection = (xmin, ymin, xmax, ymax)
    if event == cv2.EVENT_LBUTTONUP:
        drag_start = None
        track_window = selection
        selection = None


def main():
    track_window1 = ()
    cv2.namedWindow('image', 1)
    cv2.setMouseCallback('image', onmouse)
    # We will track the frames as we load them off of disk
    # for k, f in enumerate(sorted(glob.glob(os.path.join(video_folder, "*.jpg")))):
    k = 0
    while (1):
        ret, frame = cap.read()
        if not ret:
            print("Game over!")
            break
        print("Processing Frame {}".format(k))
        img_raw = frame  # cv2.imread(f)
        image = img_raw.copy()

        # We need to initialize the tracker on the first frame
        if k == 0:
            # Start a track on the object you want. box the object using the mouse and press 'Enter' to start tracking
            while True:
                ret, image = cap.read()
                img_first = image.copy()
                if track_window:
                    cv2.rectangle(img_first, (track_window[0], track_window[1]), (track_window[2], track_window[3]),
                                  (0, 0, 255), 1)
                elif selection:
                    cv2.rectangle(img_first, (selection[0], selection[1]), (selection[2], selection[3]), (0, 0, 255), 1)
                if track_window1:
                    cv2.rectangle(img_first, (track_window1[0], track_window1[1]), (track_window1[2], track_window1[3]),
                                  (0, 255, 255), 1)

                cv2.imshow('image', img_first)
                if cv2.waitKey(10) == 32:  # space
                    track_window1 = list(track_window)
                    track_window1[2] = track_window1[2] - track_window1[0]
                    track_window1[3] = track_window1[3] - track_window1[1]
                    track_window1 = tuple(track_window1)
                    k = 2
                    tracker1.init(image, track_window1)
                    break


        else:
            tracker(cap, frame, track_window1)

    cv2.destroyAllWindows()







def tracker(cam, frame, bbox):
    tracker = KCFTracker(True, True, True) # (hog, fixed_Window, multi_scale)
    tracker.init(bbox, frame)
    
    while True:
        ok, frame = cam.read()

        timer = cv2.getTickCount()
        bbox = tracker.update(frame)
        bbox = list(map(int, bbox))
        fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)

        # Tracking success
        p1 = (int(bbox[0]), int(bbox[1]))
        p2 = (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3]))
        cv2.rectangle(frame, p1, p2, (255, 0, 0), 2, 1)
        tld_roi.x_offset = int(bbox[0])
        tld_roi.y_offset = int(bbox[1])
        tld_roi.width = int(bbox[2])
        tld_roi.height = int(bbox[3])
                    # 发布ROI
        pub.publish(tld_roi)
        
        #print("左上角坐标为",p1,"右下角坐标为",p2)
        # Put FPS
        cv2.putText(frame, "FPS : " + str(int(fps)), (100, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)

        cv2.imshow("Tracking", frame)
        out.write(frame)
        # Exit if ESC pressed
        k = cv2.waitKey(1) & 0xff
        if k == 27:
            break

    while True:
        ok, frame = cam.read()

        tld_roi.x_offset = 0
        tld_roi.y_offset = 0
        tld_roi.width = 0
        tld_roi.height = 0
                    # 发布ROI
        pub.publish(tld_roi)
        
        cv2.imshow("Tracking", frame)
        out.write(frame)
        # Exit if ESC pressed
        k = cv2.waitKey(1) & 0xff
        if k == 27:
            break

    cam.release()
    cv2.destroyAllWindows()

if __name__ == '__main__':
    rospy.init_node("DSSTtracker")
    pub = rospy.Publisher("roi",ROI,queue_size=10)
    tld_roi = ROI()
    main()




